import matplotlib.pyplot as plt
import numpy as np
from scipy.interpolate import LinearNDInterpolator
from astropy.modeling import models, fitting

from scipy.ndimage import gaussian_filter
from matplotlib.colors import LogNorm

import JupiterMag as jm 


def cart2pol(x,y):
    rho = np.sqrt(x**2 + y ** 2)
    phi = np.arctan2(y,x)
    return(rho,phi)

def pol2cart(rho,phi):
    x = rho * np.cos(phi)
    y = rho * np.sin(phi)
    return(x,y)


#this is to take the corotation model with sorba and then look for the differences with a 'neutral average'

   


def Bmag(longc,latgc,model='jrm33',radial=True,R=1.00,MaxDeg=30):

    
    longc2 = 360-longc
    latgc2 = latgc+90
    
    
    longcr = longc2*np.pi/180.0
    latgcr = (180-latgc2)*np.pi/180.0
    r = np.zeros(longcr.shape) + R
    
    jm.Internal.Config(Model=model,CartesianIn=False,CartesianOut=False,Degree=MaxDeg)
    BrI,BtI,BpI = jm.Internal.Field(r,latgcr,longcr,MaxDeg=MaxDeg)
        
    B = np.sqrt(BtI**2 + BpI**2) #+ np.sqrt(BrE**2 + BtE**2 + BpE**2)
    
    #convert to Gauss
    Br = BrI #+ BrE
    if radial:
        Bg = Br.reshape(longcr.shape)*1e-5
    else:
        Bg = B.reshape(longcr.shape)*1e-5

    return Bg

# modelnames=['isaac','gsfc13ev','gsfc15ev','jpl15ev','jpl15evs','o4','o6','p11a','sha','u17ev','v117ev']
modelnames=['cassini3','cassini5','cassini11']
# |"isaac"|"gsfc13ev"|"gsfc15ev"|"gsfc15evs"|"jpl15ev"|"jpl15evs"|"o4"|"o6"|"p11a"|"sha"|"u17ev"|"v117ev"|"none" 




# %%

####  set up a scaling set of latitudes across the array


##### a quick look at the data suggest 34 pixels between  the main emission at ~12 colatitude  either side of the pole

# scaling_width = 80
scaling_width = 90

#84 = 36 pixels apart, needs to be 34
#80 = 34 pixels apart, same as auroral emission at 12 colatitude

pole_lat = np.zeros([scaling_width*2+1,scaling_width*2+1])


rotvel=(np.arange(scaling_width*2+1)-scaling_width)/scaling_width*(-9.87)
plt.figure()
plt.plot(rotvel)

for ii in range (scaling_width*2+1):
    for jj in range (scaling_width*2+1):
        i = scaling_width - ii
        j = scaling_width - jj
        if np.sqrt(i**2 + j**2) < scaling_width:
            ratio = np.sqrt(i**2 + j**2) / scaling_width
            pole_lat[jj,ii]=np.degrees(np.arccos(ratio))
        

pole_lat2=pole_lat+0
pole_lat2[pole_lat2 < 78]=0
pole_lat2[pole_lat < 77]=pole_lat[pole_lat < 77]

pl2=pole_lat2[int(scaling_width*0.75):int(scaling_width*1.25),int(scaling_width*0.75):int(scaling_width*1.25)]
pl3=pole_lat2[scaling_width,int(scaling_width*0.75):int(scaling_width*1.25)]


# plt.figure()
# # plt.imshow(pl2)
# plt.plot(pole_lat2[scaling_width,int(scaling_width*0.75):int(scaling_width*1.25)])


####  set up a scaling set of longitudes across the array



pole_lon = np.zeros([scaling_width*2+1,scaling_width*2+1])

for ii in range (scaling_width*2+1):
    for jj in range (scaling_width*2+1):
        i = scaling_width - ii
        j = scaling_width - jj
        rho,phi = cart2pol(i,j)

        if np.sqrt(i**2 + j**2) < scaling_width:
            pole_lon[jj,ii]= (360-np.mod((np.degrees(phi)+270) ,360))/360*24
        



pole_lon = np.rot90(np.rot90(pole_lon))

# pole_lon2=pole_lon+0
# pole_lon2[pole_lon2 < 76]=0
# pole_lon2[pole_lon < 74]=pole_lon[pole_lon < 74]


plt.imshow(pole_lat)
plt.show()
plt.imshow(pole_lon)
plt.show()


# %%
##### load and set up the actual lat and lon values

# nc_lat = np.load('lat.npy')

# nc_lat = np.flip(nc_lat[152:272,50:250])

# plt.imshow(nc_lat)

# nc_lon = np.load('lon.npy')

# nc_lon = (360-np.flip(nc_lon[152:272,50:250]))/360*24


nc_lat = pole_lat
nc_lon = pole_lon


fake_lon = nc_lon + 0
fake_lon[fake_lon > 12] = 24 - fake_lon[fake_lon > 12]


idx = np.unravel_index(np.argmax(nc_lat), nc_lat.shape)

pole_pos_x = idx[1]
pole_pos_y = idx[0]

#  size of equator in 0.1" pixels

# 01 aug 2017
S_sel = 31.8
S_diameter = 17.8



plt.imshow(nc_lon,origin='lower')

# plt.figure()
# plt.plot(nc_lon[60,:])



scaling_vel = (-9.87)/(S_diameter*10/2)

# rotvel=np.zeros_like(nc_lon[60,:])
# for i in range(rotvel.size):
#     rotvel[i]=(i-rotvel.size/2)*scaling_vel

# plt.plot(rotvel)



pole_lat = nc_lat
pole_lon = nc_lon




# %%

modelname=modelnames[2]

BgH = Bmag(nc_lon,nc_lat,model=modelname,radial=False,R=1)


#get the scale
            
  

# BgV = Bmag(longc,latgc,model=modelname,radial=True,MaxDeg=maxorders)
BgV = Bmag(nc_lon,nc_lat,model=modelname,radial=True,R=1)

BgA = np.degrees(np.arctan(BgV/BgH))

BgT = np.sqrt(BgV**2+BgH**2)*1e-4 # Tesla




# %%

# ####  set up a scaling set of latitudes across the array


# ##### a quick look at the data suggest 34 pixels between  the main emission at ~12 colatitude  either side of the pole

# # scaling_width = 80
# scaling_width = 90

# #84 = 36 pixels apart, needs to be 34
# #80 = 34 pixels apart, same as auroral emission at 12 colatitude

# pole_lat = np.zeros([scaling_width*2+1,scaling_width*2+1])


# rotvel=(np.arange(scaling_width*2+1)-scaling_width)/scaling_width*(-9.87)
# plt.figure()
# plt.plot(rotvel)

# for ii in range (scaling_width*2+1):
#     for jj in range (scaling_width*2+1):
#         i = scaling_width - ii
#         j = scaling_width - jj
#         if np.sqrt(i**2 + j**2) < scaling_width:
#             ratio = np.sqrt(i**2 + j**2) / scaling_width
#             pole_lat[jj,ii]=np.degrees(np.arccos(ratio))
        

# pole_lat2=pole_lat+0
# pole_lat2[pole_lat2 < 78]=0
# pole_lat2[pole_lat < 77]=pole_lat[pole_lat < 77]

# pl2=pole_lat2[int(scaling_width*0.75):int(scaling_width*1.25),int(scaling_width*0.75):int(scaling_width*1.25)]
# pl3=pole_lat2[scaling_width,int(scaling_width*0.75):int(scaling_width*1.25)]


# # plt.figure()
# # # plt.imshow(pl2)
# # plt.plot(pole_lat2[scaling_width,int(scaling_width*0.75):int(scaling_width*1.25)])


# ####  set up a scaling set of longitudes across the array



# pole_lon = np.zeros([scaling_width*2+1,scaling_width*2+1])

# for ii in range (scaling_width*2+1):
#     for jj in range (scaling_width*2+1):
#         i = scaling_width - ii
#         j = scaling_width - jj
#         rho,phi = cart2pol(i,j)

#         if np.sqrt(i**2 + j**2) < scaling_width:
#             pole_lon[jj,ii]= (360-np.mod((np.degrees(phi)+270) ,360))/360*24
        




# # pole_lon2=pole_lon+0
# # pole_lon2[pole_lon2 < 76]=0
# # pole_lon2[pole_lon < 74]=pole_lon[pole_lon < 74]



# ### use Sorba to map from ionosphere to magnetosphere R across local times 

South = "False"
Storm = "True"


if Storm == "True":
    storm_str='_comp'
else:
    storm_str='_expand'

if South == "True":
    noon_filename='Snoon_L_colat'+storm_str+'.txt'
    dawn_filename='Sdawn_L_colat'+storm_str+'.txt'
    dusk_filename='Sdusk_L_colat'+storm_str+'.txt'
    night_filename='Snight_L_colat'+storm_str+'.txt'
    lat_min = 58
    south_str='_south'
else:
    noon_filename='Nnoon_L_colat'+storm_str+'.txt'
    dawn_filename='Ndawn_L_colat'+storm_str+'.txt'
    dusk_filename='Ndusk_L_colat'+storm_str+'.txt'
    night_filename='Nnight_L_colat'+storm_str+'.txt'
    lat_min=62.6
    south_str='_north'

noon=np.loadtxt(noon_filename)

dawn=np.loadtxt(dawn_filename)

dusk=np.loadtxt(dusk_filename)

night=np.loadtxt(night_filename)





fig = plt.figure()
# fig, axes = plt.subplots(nrows=2, ncols=3)

fig.set_size_inches(8, 6.5)
           
         
# ax = axes.ravel()


noon_lat=noon[:,1]
noon_long=np.zeros_like(noon_lat)+12
noon_R=noon[:,0]


dawn_lat=dawn[:,1]
dawn_long=np.zeros_like(dawn_lat)+6
dawn_R=dawn[:,0]

night_lat=night[:,1]
night_long=np.zeros_like(night_lat)
night_R=night[:,0]

dusk_lat=dusk[:,1]
dusk_long=np.zeros_like(dusk_lat)+18
dusk_R=dusk[:,0]


x = np.hstack([night_long,0,dawn_long,6,noon_long,12,dusk_long,18,night_long+24,24])
y = np.hstack([night_lat,90,dawn_lat,90,noon_lat,90,dusk_lat,90,night_lat,90])
z = np.hstack([night_R,1,dawn_R,1,noon_R,1,dusk_R,1,night_R,1])


X = np.linspace(min(x), max(x))
Y = np.linspace(min(y), max(y))

X2 = pole_lon
Y2= 90-pole_lat


X, Y = np.meshgrid(X, Y)  # 2D grid for interpolation

interp = LinearNDInterpolator(list(zip(x, y)), z)


Z = interp(X2, Y2)

Zplot = Z+0
Zplot[nc_lat < 1] = np.nan

yvalue=1.5

levels = np.arange(0,90,10)
levels2 = np.arange(0,24,2)

# ax1.pcolormesh(X2, Y2, Z, shading='auto')
# plot_mapping=ax1.imshow(Zplot,origin='lower',extent=[-10,+10,-yvalue,+yvalue], norm=LogNorm(vmin=1, vmax=50))
# ax1.contour(nc_lat,levels,colors='grey',linestyles='dotted',extent=[-10,+10,-yvalue,+yvalue])
# ax1.contour(fake_lon,levels2,colors='grey',linestyles='dotted',extent=[-10,+10,-yvalue,+yvalue])
# ax1.plot([0,0],[-1,1.49],color='grey',linestyle='dotted')
# mcb=plt.colorbar(plot_mapping,ax=ax1,label='Magnetospheric equatorial R$_S$', orientation="horizontal", ticks=[1, 10, 50])
# mcb.ax.set_xticklabels(['1', '10', '50']) 

# ax1.set_xlim([-7,7])

# %%




imw_r15_filename='imw_subcorot_r15.txt'
imw_r15=np.loadtxt(imw_r15_filename)

r15_pos=np.hstack([imw_r15[:,0],55,50])
r15_sub=np.hstack([imw_r15[:,1],0.99,1])


# plt.plot(r15_pos,r15_sub)
lat_min_b=52

difflatmin2=lat_min_b-2


xx=np.arange(91-difflatmin2)+difflatmin2

g_init = models.Polynomial1D(degree=18) #,c0=c0,c1=c1,c2=c2
#g_init = models.Gaussian1D()+ models.Polynomial1D(degree=2)
fit_g = fitting.LevMarLSQFitter()
g = fit_g(g_init, r15_pos, r15_sub,maxiter=1000)

gxx = g(xx)
gxx[xx<60.6]=0.96+(62-xx[xx<60.6])/300
# gxx[xx<66]=0.9
gxx[xx>88]=0.52

# plt.plot(xx,gxx)




imw_neutral_subcorot_r15 = g(nc_lat)

imw_neutral_subcorot_r15[nc_lat < 60.6] = 0.96+(62-nc_lat[nc_lat < 60.6])/300
imw_neutral_subcorot_r15[nc_lat >88] = 0.52




imw_r6_filename='imw_subcorot_r6.txt'
imw_r6=np.loadtxt(imw_r6_filename)

r6_pos=np.hstack([imw_r6[:,0],55,50])
r6_sub=np.hstack([imw_r6[:,1],0.99,1])


# plt.plot(r6_pos,r6_sub)
lat_min_b=52

difflatmin2=lat_min_b-2


xx=np.arange(91-difflatmin2)+difflatmin2

g_init = models.Polynomial1D(degree=18) #,c0=c0,c1=c1,c2=c2
#g_init = models.Gaussian1D()+ models.Polynomial1D(degree=2)
fit_g = fitting.LevMarLSQFitter()
g = fit_g(g_init, r6_pos, r6_sub,maxiter=1000)

gxx = g(xx)
gxx[xx<59]=0.96+(62-xx[xx<59])/260
gxx[xx<52]=1
# gxx[xx<66]=0.9
gxx[xx>88]=0.765

# plt.plot(xx,gxx)



imw_neutral_subcorot_r6 = g(nc_lat)

imw_neutral_subcorot_r6[nc_lat < 59] = 0.96+(62-nc_lat[nc_lat < 59])/360
imw_neutral_subcorot_r6[nc_lat < 52] = 1
imw_neutral_subcorot_r6[nc_lat >88] = 0.765




# %%

#######  Use wilson to plot from R to a subrotation rate 

wilson_filename='wilson2.txt'
wilson=np.loadtxt(wilson_filename)


# calculate the corotation speed

# 40 x 4s and 40 x 30 making the start and end limits very hard to avoid for the polynomial 
x = np.hstack([4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,4,wilson[:,0],31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31,31])
y = np.hstack([1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,wilson[:,1]/(wilson[:,0]*9.87),0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333,0.333])

xx=np.arange(34)+1

g_init = models.Polynomial1D(degree=4) #,c0=c0,c1=c1,c2=c2
#g_init = models.Gaussian1D()+ models.Polynomial1D(degree=2)
fit_g = fitting.LevMarLSQFitter()
g = fit_g(g_init, x, y,maxiter=1000)

gxx = g(xx)
gxx[xx<4]=1
gxx[xx>30]=0.333

print(g.param_names)


    # Plot the data with the best-fit model
    # plt.figure(figsize=(8,5))
# ax4.plot(x, y, marker=".",linestyle='None')
# ax4.plot(xx, gxx,c='darkred')
# ax4.set_xlabel('Magnetospheric equatorial R$_S$')
# ax4.set_ylabel('Sub-corotation')

# ax5.xlabel('Position')
# ax5.ylabel('Flux')







# cov = fit_g.fit_info['param_cov']


# ax5.set_title("Downscaled image (no aliasing)")

# ax1.set_xlim(0, 512)
# ax1.set_ylim(512, 0)













#######  Map from Sorba to Wilson to map subrotation rate across the planet 

# firstly, if mapping is Nan, that means you are in the polar cap, so use a subcorotation of 0.3 (or some other escape clause)

bbox_props2_1 = dict(boxstyle="circle,pad=0.05", fc="antiquewhite", ec="sandybrown", lw=2)
bbox_props3_1 = dict(boxstyle="round,pad=0.3", fc="antiquewhite", ec="grey", lw=2)

ax0 = fig.add_axes([0., 0.5, 1/3, 1/2-0.02]) # main axes


plot_mapping=ax0.imshow(Zplot,origin='lower',norm=LogNorm(vmin=1, vmax=50))
ax0.contour(nc_lat,levels,colors='grey',linestyles='dotted')
ax0.contour(fake_lon,levels2,colors='grey',linestyles='dotted')
ax0.plot([90,90],[0,90],color='black',linestyle='dashdot')
mcb=plt.colorbar(plot_mapping,ax=ax0,label='Magnetospheric equatorial R$_S$', orientation="horizontal", ticks=[1, 10, 50],pad=0.05,shrink=0.9)
mcb.ax.set_xticklabels(['1', '10', '50']) 
ax0.set_axis_off()
# ax0.set_xlim([-7,7])
ax0.text(158,158,'a',bbox=bbox_props2_1)
ax0.text(90,20,'12UT',bbox=bbox_props3_1,ha='center', va='center')


ax3 = fig.add_axes([(1/3*0.2), 0+0.05, 1/3-(1/3*0.3), 1/2*0.86-0.05]) # main axes

ax3.plot(x, y, marker=".",linestyle='None')
ax3.plot(xx, gxx,c='darkred')
ax3.set_xlabel('Magnetospheric equatorial R$_S$')
ax3.set_ylabel('Sub-corotation')
ax3.text(30,0.95,'b',bbox=bbox_props2_1)


subrot=g(Z)

subrot[Z < 4]=1
subrot[Z > 30]=0.3333

# subrot[np.isnan(subrot)]=0.3333
subrot[np.isnan(subrot)]=0.3

losvel=np.zeros_like(subrot)

subrotplot = subrot+0
subrotplot[nc_lat < 1] = np.nan



ax1 = fig.add_axes([1/3, 1/2, 1/3, 1/2-0.02]) # main axes

# ax4.plot(Z)
plot_subrot=ax1.imshow(subrotplot,origin='lower',vmin=0,vmax=1,cmap='magma')
# ax1.set_ylim([0,15])
ax1.contour(nc_lat,levels,colors='grey',linestyles='dotted')
ax1.contour(fake_lon,levels2,colors='grey',linestyles='dotted')
ax1.plot([90,90],[0,90],color='black',linestyle='dashdot')
plt.colorbar(plot_subrot,ax=ax1,label='Ion Sub-corotation', orientation="horizontal",pad=0.05,shrink=0.9)
# ax1.set_xlim([-4,4])
ax1.axis('off')
ax1.text(158,158,'c',bbox=bbox_props2_1)
ax1.text(90,20,'12UT',bbox=bbox_props3_1,ha='center', va='center')

# for i in range(subrot[:,0].size):
#     losvel[i,:]=subrot[i,:]*rotvel

seeing=0.67
losvelseeing = gaussian_filter(losvel, sigma=seeing*10/2.355)

# 0.5" is 5 pixel - but is FWHM - convert to sigma fwhm = 2.355 sigma

 
# ax5.imshow(losvel)
arcsec_x = np.arange(200)*0.1-10

# ax5.plot(arcsec_x,losvel[pole_pos_y,:])
# ax5.plot(arcsec_x,losvelseeing[pole_pos_y,:])
# ax5.set_xlim([-7,7])
# ax5.set_ylim([-4,4])
# ax5.set_xlabel('arcsec (x)')
# ax5.set_ylabel('Velocity (km/s)')


# velplot4=ax5.imshow(losvelseeing,origin='lower',extent=[-10,+10,-yvalue,+yvalue],vmin=-3,vmax=3,cmap='seismic_r')
# mcb=plt.colorbar(velplot4,ax=ax5,label='Mean Sub-corotation component', orientation="horizontal",pad=0.05,shrink=0.9)

# ax5.set_xlabel('arcsec (x)')
# ax5.set_ylabel('arcsec (y)')
# # ax5.contour(Xplot,[0.5,1.5],colors='orange',extent=[-10,+10,-yvalue,+yvalue])
# # ax5.contour(Zplot,[0,18,100],colors='blue',extent=[-10,+10,-yvalue,+yvalue])

# ax5.set_ylim([-1.5,1.5])
# ax5.set_xlim([-7,7])


plt.tight_layout()
# plt.subplots_adjust(bottom=0.0, right=1, top=1)

difflatmin=lat_min-2

# np.save('losvel.npy',losvel)
# np.save('losvelseeing.npy',losvelseeing)

ncl2=nc_lat[nc_lat >difflatmin]
srp2=subrotplot[nc_lat >difflatmin]



# ax4 = fig.add_axes([1/3+(1/3*0.2), 0, 1/3-(1/3*0.3), 1/2*0.86]) # main axes
ax4= fig.add_axes([1/3+(1/3*0.2), 0+0.05, 1/3-(1/3*0.3), 1/2*0.86-0.05]) # main axes

ax4.plot(ncl2,srp2, marker=",",linestyle='None')




xx=np.arange(91-difflatmin)+difflatmin

g_init = models.Polynomial1D(degree=18) #,c0=c0,c1=c1,c2=c2
# g_init = models.Polynomial1D(degree=25) #,c0=c0,c1=c1,c2=c2
#g_init = models.Gaussian1D()+ models.Polynomial1D(degree=2)
fit_g = fitting.LevMarLSQFitter()
g = fit_g(g_init, ncl2, srp2,maxiter=1000)

gxx = g(xx)
gxx[xx<60]=1
# gxx[xx<66]=0.9
gxx[xx>85]=0.3

meanvel = np.zeros((60))

# print(nc_lat[np.where(np.logical_and(nc_lat>= latvalue/2, nc_lat<= latvalue/2+0.5))])
for latvalue in range(60*2,90*2): 
    latmask = np.where(np.logical_and(nc_lat>= latvalue/2, nc_lat<= latvalue/2+0.5))
    meanvel[latvalue-120]=np.nanmean(subrotplot[latmask])

ax4.plot(np.arange(60,90,0.5)+0.25, meanvel,c='red',marker='x',linestyle='none',markersize=5)

# ax4.plot(x, y, marker=".",linestyle='None')
ax4.plot(xx, gxx,c='darkred',alpha=0.8)
ax4.set_xlabel('Northern latitude')
ax4.set_ylabel('Ion Sub-corotation with latitude')
ax4.text(88,0.95,'d',bbox=bbox_props2_1)


neutral_subcorot = g(nc_lat)

neutral_subcorot[nc_lat < difflatmin] = subrotplot[nc_lat < difflatmin]
neutral_subcorot[nc_lat >85] = 0.3



veldiff = subrotplot - neutral_subcorot
veldiff_r15 = subrotplot - imw_neutral_subcorot_r15
veldiff_r6 = subrotplot - imw_neutral_subcorot_r6

# veldiff = subrotplot - imw_neutral_subcorot_r6

Efield = -veldiff*BgT*1000


# veldiff[nc_lat < 65] = np.nan

ax2 = fig.add_axes([2/3, 1/2, 1/3, 1/2-0.02]) # main axes

nscr=ax2.imshow(neutral_subcorot,origin='lower',vmin=0,vmax=1,cmap='magma')
# ax2.set_xlim([-4,4])
plt.colorbar(nscr,ax=ax2,label='Neutral Sub-corotation', orientation="horizontal",pad=0.05,shrink=0.9)
ax2.contour(nc_lat,levels,colors='grey',linestyles='dotted')
ax2.contour(fake_lon,levels2,colors='grey',linestyles='dotted')
ax2.plot([90,90],[0,90],color='black',linestyle='dashdot')
ax2.set_axis_off()
ax2.text(158,158,'e',bbox=bbox_props2_1)
ax2.text(90,20,'12UT',bbox=bbox_props3_1,ha='center', va='center')

# ax2.set_ylim([-2,2])
# ax2.set_axis_off()
# for tick in ax2.get_xaxis().get_major_ticks():
#     tick.set_pad(0.)

ax5 = fig.add_axes([2/3, 0, 1/3, 1/2-0.02]) # main axes

plot_diffvel=ax5.imshow(veldiff,origin='lower',cmap='RdBu')
plt.colorbar(plot_diffvel,ax=ax5,label='$\Delta$ Sub-corotation', orientation="horizontal",pad=0.05,shrink=0.9)
ax5.contour(nc_lat,levels,colors='grey',linestyles='dotted')
ax5.contour(fake_lon,levels2,colors='grey',linestyles='dotted')
ax5.plot([90,90],[0,90],color='black',linestyle='dashdot')
# ax5.set_xlim([-4,4])
ax5.set_axis_off()
ax5.text(158,158,'f',bbox=bbox_props2_1)
ax5.text(90,20,'12UT',bbox=bbox_props3_1,ha='center', va='center')

ax5.contour(Zplot,[0,8,300],colors='green',linestyles='dotted')
ax5.contour(Zplot,[0,10,500],colors='green',linestyles='dotted')

Zplotmask=Zplot*1.
Zplotmask[0:50,:]=0.
Zplotmask[-50:,:]=0.
Zplotmask[:,-50:]=0.
Zplotmask[:,0:50]=0.
Zplotmask[np.isfinite(Zplotmask)]=1.
Zplotmask[~np.isfinite(Zplotmask)]=0.


ax5.contour(Zplotmask,[0,1,300],colors='darkgreen',linestyles='dotted')


np.save('model_vels'+south_str+storm_str+'.npy',veldiff)
np.save('model_vels'+south_str+storm_str+'_r15.npy',veldiff_r15)
np.save('model_vels'+south_str+storm_str+'_r6.npy',veldiff_r6)
np.save('model_dist'+south_str+storm_str+'.npy',Zplot)

# ax5.set_xticks([])
# ax5.set_yticks([])
# ax5.set_position([2/3,0,1/3,0.5])
# plt.tight_layout()
plt.savefig('asd_diff_corot_model_pole'+south_str+storm_str+'.pdf', dpi=300, bbox_inches='tight', facecolor='white', pad_inches=0) 
plt.show()
# 


